﻿#pragma once
#include<iostream>
using namespace std;

enum Colour
{
	red,
	black
};

template<class K, class V>
struct RBTreeNode
{
	// 这⾥更新控制平衡也要加⼊parent指针
	pair<K, V> _kv;
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	Colour _col;
	RBTreeNode(const pair<K, V>& kv)
		:_kv(kv)
		,_left(nullptr)
		,_right(nullptr)
		,_parent(nullptr)
		,_col(red)
	{}

};

template<class K, class V>
class RBTree
{
private:
	using Node = RBTreeNode<K, V>;
public:
	RBTree(const pair<K, V>& kv)
	{
		_root = new Node(kv);
	}
	RBTree() = default;

	//插入
	bool Insert(const pair<K, V>& data)
	{
		if (_root == nullptr)
		{
			_root = new Node(data);
			_root->_col = black;
			return true;
		}

		//正常插入
		Node* cur = _root;
		Node* cur_parent = nullptr;
		while (cur != nullptr)
		{
			if (data.first > cur->_kv.first)
			{
				cur_parent = cur;
				cur = cur->_right;
			}
			else if (data.first < cur->_kv.first)
			{
				cur_parent = cur;
				cur = cur->_left;
			}
			else
				return false;
		}
		if (data.first < cur_parent->_kv.first)
		{
			cur_parent->_left = new Node(data);
			cur = cur_parent->_left;
			cur->_parent = cur_parent;
		}
		else
		{
			cur_parent->_right = new Node(data);
			cur = cur_parent->_right;
			cur->_parent = cur_parent;
		}

		//当父节点为红的情况
		if (cur_parent->_col == red)
		{
			while (cur_parent && cur_parent->_col != black)
			{
				//调整父、叔、爷结点
				Node* cur_grandfather = cur_parent->_parent;
				Node* cur_uncle = cur_grandfather->_left;
				if (cur_parent == cur_uncle)
					cur_uncle = cur_grandfather->_right;

				//叔结点存在且为红
				if (cur_uncle && cur_uncle->_col == red)
				{
					cur_grandfather->_col = red;
					cur_parent->_col = cur_uncle->_col = black;
				}

				//叔结点不存在或存在为黑
				else
				{
					if (cur_grandfather->_left == cur_parent)
					{
						if (cur_parent->_left == cur)
						{
							RotateR(cur_grandfather);
							cur_parent->_col = black;
							cur_grandfather->_col = red;
						}
						else
						{
							RotateL(cur_parent);
							RotateR(cur_grandfather);
							cur->_col = black;
							cur_grandfather->_col = red;
						}

					}
					else
					{
						if (cur_parent->_right == cur)
						{
							RotateL(cur_grandfather);
							cur_parent->_col = black;
							cur_grandfather->_col = red;
						}
						else
						{
							RotateR(cur_parent);
							RotateL(cur_grandfather);
							cur->_col = black;
							cur_grandfather->_col = red;
						}
					}
				}

				cur = cur_grandfather;
				cur_parent = cur->_parent;
			}
			_root->_col = black;	
		}
		return true;
	}
		//右旋
	void RotateR(Node* head)
	{
		Node* head_parent = head->_parent;
		Node* head_left = head->_left;
		Node* head_left_right = head->_left->_right;

		head->_left = head_left_right;
		if(head_left_right)                                
			head_left_right->_parent = head;

		head_left->_right = head;
		head->_parent = head_left;

		if (head_parent)      
		{
			head_left->_parent = head_parent;
			if (head_left->_kv.first < head_parent->_kv.first)
				head_parent->_left = head_left;
			else
				head_parent->_right = head_left;
		}
		else
		{
			_root = head_left;
			_root->_parent = nullptr;
		}

	}

	//左旋
	void RotateL(Node* head)
	{
		Node* head_parent = head->_parent;
		Node* head_right = head->_right;
		Node* head_right_left = head->_right->_left;

		head->_right = head_right_left;
		if (head_right_left)
			head_right_left->_parent = head;

		head_right->_left = head;
		head->_parent = head_right;

		if (head_parent)
		{
			head_right->_parent = head_parent;
			if (head_right->_kv.first < head_parent->_kv.first)
				head_parent->_left = head_right;
			else
				head_parent->_right = head_right;
		}
		else
		{
			_root = head_right;
			_root->_parent = nullptr;
		}
	}
		
	//查找
	Node* Find(const K& key)
	{
		Node* cur = _root;
		while (cur != nullptr)
		{
			if (key < cur->_kv.first)
				cur = cur->_left;
			else if (key > cur->_kv.first)
				cur = cur->_right;
			else
				return cur;
		}
		return nullptr;
	}

	int Size()
	{
		return _Size(_root);
	}

	int Height()
	{
		return _Height(_root);
	}

	//中序遍历
	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

	//判断是否为红黑树
	bool IsRBTree()
	{
		if(_root->_col == black)
			return IsRBTree_(_root);
		return false;
	}

private:
	int _Size(Node* root)
	{
		if (root == nullptr)
			return 0;

		int left_size = _Size(root->_left);
		int right_size = _Size(root->_right);
		
		return left_size + right_size + 1;
	}

	int _Height(Node* root)
	{
		if (root == nullptr)
			return 0;

		int left_height = _Height(root->_left);
		int right_height = _Height(root->_right);

		return left_height > right_height ? left_height + 1 : right_height + 1;
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr)
			return;

		_InOrder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << " ";
		_InOrder(root->_right);
	}

	bool condition3(Node* root)
	{
		if (root == nullptr)
			return true;

		if (root->_col == red)
		{
			if (root->_parent && root->_parent->_col == red)
				return false;
		}
		if (condition3(root->_left))
		{
			if (condition3(root->_right))
				return true;
		}
		return false;
	}
	bool _IsRBTree(Node* root, int black_record, int& record)
	{
		if (root == nullptr)
		{
			if (black_record == record)
				return true;
			else
				return false;
		}

		if (root->_col == black)
			++black_record;
		if (_IsRBTree(root->_left, black_record, record))
		{
			if (_IsRBTree(root->_right, black_record, record))
				return true;
		}
		return false;
	}
	bool IsRBTree_(Node* root)
	{
		if (condition3(root))
		{
			int record = 0;
			Node* cur = root;
			while (cur != nullptr)
			{
				if (cur->_col == black)
					++record;
				cur = cur->_left;
			}
			return _IsRBTree(root, 0, record);
		}
		return false;
	}

	Node* _root = nullptr;
};



